# === 1. 基础依赖 =============================================================
import sys
import cv2
import requests
import numpy as np
from collections import Counter, deque
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtCore import QThread, QTimer, Qt
from PyQt5.QtWidgets import QFileDialog, QMessageBox
import os

# 读取服务器地址
from main import load_settings

cfg = load_settings()
SERVER_URL = f"http://{cfg['ip']}:{cfg['port']}/detect"


# === 2. 解码线程 =============================================================
class DecoderThread(QThread):
    """
    负责把本地视频逐帧解码，放入 raw_q（双端队列）中。
    支持暂停 / 继续 / 跳转 / 重启。
    """
    frame_idx_updated = QtCore.pyqtSignal(int)  # 当前帧号

    def __init__(self, video_path, raw_q, parent=None):
        super().__init__(parent)
        self.setTerminationEnabled(True)  # 允许强制终止
        self.video_path = video_path
        self.raw_q = raw_q  # 与推理线程共享
        self._run = True  # 总开关
        self._pause = False  # 暂停标志
        self.cap = None  # cv2.VideoCapture 实例

    # ---------------- 线程主循环 ----------------
    def run(self):
        self.cap = cv2.VideoCapture(self.video_path)
        if not self.cap.isOpened():
            return
        fps = self.cap.get(cv2.CAP_PROP_FPS) or 30
        delay = int(1000 / fps)  # 按帧率计算每帧间隔(ms)
        while self._run and self.cap.isOpened():
            if self._pause:  # 暂停时自旋
                self.msleep(30)
                continue
            ok, frame = self.cap.read()
            if not ok:  # 读到结尾
                break
            self.raw_q.append(frame)  # 把最新帧放到队列
            self.frame_idx_updated.emit(
                    int(self.cap.get(cv2.CAP_PROP_POS_FRAMES))
                    )
            self.msleep(delay // 2)  # 简单限速，防止过快
        self.cap.release()
        self.cap = None

    # ---------------- 控制接口 ----------------
    def pause(self):
        self._pause = True

    def resume(self):
        self._pause = False

    def stop(self):
        self._run = False
        self._pause = False

    def seek(self, frame_no):
        """跳转到指定帧"""
        if self.cap and self.cap.isOpened():
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, frame_no)

    def restart(self):
        """重新打开视频并回到第 0 帧"""
        if self.cap:
            self.cap.release()
        self.cap = cv2.VideoCapture(self.video_path)
        if self.cap.isOpened():
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        self._run = True
        self._pause = False


# === 3. 推理线程 =============================================================
class InferenceThread(QThread):
    """
    不断从 raw_q 取最新帧，调用远端推理接口，把结果通过信号发回主线程。
    """
    result_ready = QtCore.pyqtSignal(np.ndarray, np.ndarray, str)

    def __init__(self, url, raw_q, parent=None):
        super().__init__(parent)
        self.url = url
        self.raw_q = raw_q
        self._run = True
        self.session = requests.Session()  # 复用 TCP 连接
        self.setTerminationEnabled(True)

    def run(self):
        while self._run:
            if not self.raw_q:  # 队列为空
                self.msleep(5)
                continue
            frame = self.raw_q[-1]  # 只拿最新帧，避免堆积
            ok, enc = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 75])
            if not ok:
                continue
            try:
                r = self.session.post(
                        self.url,
                        files={'file': ('f.jpg', enc.tobytes(), 'image/jpeg')},
                        timeout=30000  # 30 秒超时
                        )
                if r.status_code == 200:
                    data = r.json()
                    cls_list = data.get("cls_list", [])
                    img_hex = data.get("img", "")
                    img_arr = np.frombuffer(bytes.fromhex(img_hex), np.uint8)
                    ann = cv2.imdecode(img_arr, cv2.IMREAD_COLOR)
                    counter = Counter(cls_list)
                    stat_text = "\n".join([f"{k}：{v}" for k, v in counter.items()])
                    cls_text = stat_text if stat_text else "未检测到物体"
                else:
                    cls_text = "error"
                    ann = frame
            except Exception as e:
                print('Inference error:', e)
                cls_text = "error"
                ann = frame
            self.result_ready.emit(frame, ann, cls_text)

    def stop(self):
        self._run = False


# === 4. 主窗口 ===============================================================
class MainWindow(QtWidgets.QMainWindow):

    def __init__(self):
        super().__init__()
        self.setWindowTitle("视频检测（类别统计）")
        self.resize(1700, 1000)
        self.setStyleSheet("background-color: #fffef8;")

        # ---------------- 布局 ----------------
        central = QtWidgets.QWidget(self)
        self.setCentralWidget(central)
        main_layout = QtWidgets.QHBoxLayout(central)
        main_layout.setContentsMargins(10, 10, 10, 10)
        main_layout.setSpacing(10)

        # ---------- 左侧 ----------
        left_w = QtWidgets.QWidget()
        left_w.setStyleSheet(
                """
                            background: qlineargradient(x1:0, y1:0, x2:0, y2:1,
                                                        stop:0 #c7e6d5, stop:1 #7eb394);
                            border-radius: 10px;
                        """
                )
        left_layout = QtWidgets.QVBoxLayout(left_w)
        left_layout.setContentsMargins(15, 15, 15, 15)
        left_layout.setSpacing(15)

        # 原图标签
        self.label_src = QtWidgets.QLabel("原图")
        self.label_src.setMinimumHeight(200)
        self.label_src.setAlignment(QtCore.Qt.AlignCenter)
        self.label_src.setStyleSheet(
                """
                            background-color: #ffffff; color: #5d8a6a;
                            font-family: "PingFang SC"; font-size: 20px;
                            border: 1px solid #5d8a6a; border-radius: 8px; padding: 10px;
                        """
                )
        self.label_src.setScaledContents(True)
        self.label_src.setSizePolicy(
                QtWidgets.QSizePolicy.Ignored,
                QtWidgets.QSizePolicy.Ignored
                )
        left_layout.addWidget(self.label_src, 2)

        # 按钮样式
        btn_qss = """
            QPushButton{
                background-color: #5d8a6a; color: #ffffff;
                font-family: "PingFang SC"; font-size: 18px;
                padding: 10px 20px; border: none; border-radius: 5px;
            }
            QPushButton:hover{ background-color: #4a6f55; }
        """

        self.btn_open = QtWidgets.QPushButton("打开视频")
        self.btn_open.setStyleSheet(btn_qss)
        self.btn_open.clicked.connect(self.open_video)
        left_layout.addWidget(self.btn_open)

        self.btn_pause = QtWidgets.QPushButton("暂停")
        self.btn_pause.setStyleSheet(btn_qss)
        self.btn_pause.clicked.connect(self.toggle_play_pause)
        left_layout.addWidget(self.btn_pause)

        self.btn_restart = QtWidgets.QPushButton("重新播放")
        self.btn_restart.setStyleSheet(btn_qss)
        self.btn_restart.clicked.connect(self.restart_video)
        left_layout.addWidget(self.btn_restart)

        self.btn_back = QtWidgets.QPushButton("返回")
        self.btn_back.setStyleSheet(btn_qss)
        self.btn_back.clicked.connect(self.go_to_main_1)
        left_layout.addWidget(self.btn_back)
        left_layout.addStretch()

        # ---------- 右侧 ----------
        right_w = QtWidgets.QWidget()
        right_layout = QtWidgets.QVBoxLayout(right_w)
        right_layout.setContentsMargins(0, 0, 0, 0)
        right_layout.setSpacing(5)

        # 检测结果标签
        self.label_dst = QtWidgets.QLabel("检测结果")
        self.label_dst.setAlignment(QtCore.Qt.AlignCenter)
        self.label_dst.setStyleSheet(
                """
                            background-color: #f2f5f9; color: #5d8a6a;
                            font-family: "PingFang SC"; font-size: 36px;
                            border: 2px dashed #5d8a6a; border-radius: 10px;
                        """
                )
        self.label_dst.setScaledContents(True)
        self.label_dst.setSizePolicy(
                QtWidgets.QSizePolicy.Ignored,
                QtWidgets.QSizePolicy.Ignored
                )
        right_layout.addWidget(self.label_dst, 4)

        # 类别统计标签
        self.label_class = QtWidgets.QLabel("类别统计：-")
        self.label_class.setAlignment(QtCore.Qt.AlignLeft | QtCore.Qt.AlignTop)
        self.label_class.setWordWrap(True)
        self.label_class.setStyleSheet(
                """
                            background-color: #ffffff; color: #5d8a6a;
                            font-family: "PingFang SC"; font-size: 20px;
                            border: 1px solid #5d8a6a; border-radius: 8px; padding: 10px;
                        """
                )
        right_layout.addWidget(self.label_class, 1)

        main_layout.addWidget(left_w, 1)
        main_layout.addWidget(right_w, 2)

        # ---------------- 成员变量 ----------------
        self.video_path = None
        self.total_frames = 0
        self.current_frame_idx = 0
        self.server_url = SERVER_URL

        # 线程/队列
        self.raw_q = deque(maxlen=2)  # 只保留最近 2 帧，防止内存膨胀
        self.decoder_thread = None
        self.infer_thread = None
        self.display_timer = QTimer(self)
        self.display_timer.timeout.connect(self.on_display_timeout)

    # ---------------- 功能函数 ----------------
    def open_video(self):
        """弹出文件选择框，加载视频"""
        path, _ = QFileDialog.getOpenFileName(
                self, "选择视频", "", "视频文件 (*.mp4 *.avi *.mov *.mkv)"
                )
        if not path:
            return
        cap = cv2.VideoCapture(path)
        if not cap.isOpened():
            QMessageBox.warning(self, "错误", "无法打开选中的视频文件！")
            return
        self.total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        cap.release()

        self.video_path = path
        self.reset_all_threads()
        QMessageBox.information(self, "提示", f"已加载视频：{os.path.basename(path)}")

    def toggle_play_pause(self):
        """暂停/继续切换"""
        if self.decoder_thread is None:
            QMessageBox.warning(self, "警告", "请先选择视频！")
            return
        if self.decoder_thread._pause:
            self.decoder_thread.resume()
            self.display_timer.start(30)  # 30ms 刷新一次
            self.btn_pause.setText("暂停")
        else:
            self.decoder_thread.pause()
            self.display_timer.stop()
            self.btn_pause.setText("继续")

    def restart_video(self):
        """重新播放当前视频"""
        if not self.video_path:
            QMessageBox.warning(self, "警告", "请先选择视频！")
            return
        self.reset_all_threads()
        self.decoder_thread.resume()
        self.display_timer.start(30)
        self.btn_pause.setText("暂停")

    def on_display_timeout(self):
        """定时刷新（目前空实现，可扩展）"""
        pass

    def reset_all_threads(self):
        """安全地结束并重开解码/推理线程"""
        if self.decoder_thread and self.decoder_thread.isRunning():
            self.decoder_thread.stop()
            self.decoder_thread.terminate()
            self.decoder_thread.wait(1000)
        if self.infer_thread and self.infer_thread.isRunning():
            self.infer_thread.stop()
            self.infer_thread.terminate()
            self.infer_thread.wait(1000)

        self.raw_q.clear()
        self.decoder_thread = None
        self.infer_thread = None
        self.display_timer.stop()

        if self.video_path:
            self.decoder_thread = DecoderThread(self.video_path, self.raw_q)
            self.decoder_thread.start()

            self.infer_thread = InferenceThread(self.server_url, self.raw_q)
            self.infer_thread.result_ready.connect(self.show_pair)
            self.infer_thread.start()

    def show_pair(self, raw_bgr, ann_bgr, cls_text):
        """把推理结果画到界面上"""
        rgb_raw = cv2.cvtColor(raw_bgr, cv2.COLOR_BGR2RGB)
        rgb_ann = cv2.cvtColor(ann_bgr, cv2.COLOR_BGR2RGB)
        self.display_frame(rgb_raw, self.label_src)
        self.display_frame(rgb_ann, self.label_dst)
        self.label_class.setText(f"类别统计：\n{cls_text}")

    def display_frame(self, rgb, label):
        """把 numpy RGB 图像显示到 QLabel"""
        h, w, ch = rgb.shape
        bytes_per_line = ch * w
        qimg = QtGui.QImage(rgb.data, w, h, bytes_per_line, QtGui.QImage.Format_RGB888)
        label.setPixmap(
                QtGui.QPixmap.fromImage(qimg).scaled(
                        label.size(), QtCore.Qt.KeepAspectRatio, QtCore.Qt.SmoothTransformation
                        )
                )

    def go_to_main_1(self):
        """返回主界面（main.py 中的 UI_main_1）"""
        if self.decoder_thread and self.decoder_thread.isRunning():
            self.decoder_thread.pause()
        self.display_timer.stop()
        self.btn_pause.setText("继续")
        if self.decoder_thread and self.decoder_thread.isRunning():
            self.decoder_thread.terminate()
            self.decoder_thread.wait(500)
        if self.infer_thread and self.infer_thread.isRunning():
            self.infer_thread.terminate()
            self.infer_thread.wait(500)
        self.close()
        try:
            from main import UI_main_1

            self.main_win = UI_main_1()
            self.main_win.show()
        except ImportError:
            QMessageBox.warning(self, "提示", "找不到 main.py，无法返回主界面！")


# === 5. 入口 =================================================================
if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    win = MainWindow()
    win.show()
    sys.exit(app.exec_())
